#include "CNN/ckks_mnist_helper.hpp"
#include "CNN/mnist_helper.hpp"
#include "intel_itt_wrapper.hpp"

// Where to find the MNIST dataset.
std::string kDataRoot = "build/data";

auto main(int argc, char **argv) -> int {
  INTEL_ITT_PAUSE; //! ITT - Stop

  torch::manual_seed(1);
  unsigned long test_num = 1;

  // Defualt BGV Params Set
  params<helib::CKKS> params;

  // Create a `command_line_parser` using `params<helib::BGV>`
  command_line_parser parser(params);
  parser.add_option("-tn", "2", "test_num", &test_num);

  //* Parse command line
  int ret = parser.parse(argc, argv);
  if (ret) {
    return ret;
  }

  // Create a cryptograph client
  ckks_hmnist_client client(params);

  // Create a HMM status instance
  hmm_status<shmm_engine> hmm_status(METHOD::NAIVE, omp_get_max_threads());
  ckks_hmnist_server server(client.get_context(), hmm_status);

  torch::Device device(torch::kCPU);
  Net model(client, server);
  model.to(device);

  std::string pt_file = kDataRoot + "/mnist.pt";
  load_model(model, device, pt_file);

  test(model, device, kDataRoot, test_num);
}
